import os

import click
import dotenv
import openai
import pandas as pd
from tqdm import tqdm
from tree_rag.agents import DEFAULT_MODEL_NAME

dotenv.load_dotenv()

client = openai.OpenAI(
    base_url=os.environ["OPENAI_API_BASE"],
    api_key=os.environ["OPENAI_API_KEY"],
)

DEFAULT_MODEL_NAME = os.getenv("MODEL_NAME")


@click.command()
@click.option("--input_file", type=str, help="输入文件")
def main(input_file: str):
    current_dir = os.path.dirname(os.path.abspath(__file__))

    df = pd.read_excel(input_file)

    answer_list = df["answer"].tolist()

    PROMPT_TEMPLATE = """
    知识：
    {answer}

    我想创建一个QA对，请根据知识，提取出问题。直接返回问题，不要包含其他内容。
    """

    question_list = []
    for answer in tqdm(answer_list):
        messages = []
        messages.append(
            {"role": "user", "content": PROMPT_TEMPLATE.format(answer=answer)}
        )

        chat_completion = client.chat.completions.create(
            messages=messages,
            model=DEFAULT_MODEL_NAME,
        )
        response_str = chat_completion.choices[0].message.content
        question_list.append(response_str)

    df["question"] = question_list
    df.to_excel(input_file, index=False)


if __name__ == "__main__":
    main()